# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/3 11:16 
@Author  : ZhangShenao 
@File    : 6.消息的过滤与裁剪.py 
@Desc    : 消息的过滤与裁剪

过滤和裁剪消息时,一般不会修改整个Graph的状态
而是在调用LLM时,只传递特定条数的消息或者按照token长度进行修剪
通常将消息过滤封装成一个单独的函数,在调用LLM前执行该函数,而不是将消息过滤定义成一个Graph节点
这样在使用LLM时就可以避免全部将消息传递过去
并且整个Graph架构内,状态仍然会保存最完整的信息

在LangGraph中封装了trim_messages函数
可以基于多种策略对消息进行过滤和裁剪
并且trim_messages函数使用@_runnable_support装饰器进行了装饰
因此它也是一个Runnable可运行组件,可以直接拼接到LCEL表达式构建的链应用中
"""
import os

import dotenv
from langchain_core.messages import HumanMessage, AIMessage, trim_messages
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载环境变量
dotenv.load_dotenv()

# 设置消息列表
messages = [
    HumanMessage(content="你好，我叫zsa，我喜欢游泳打篮球，你喜欢什么呢？"),
    AIMessage([
        {"type": "text", "text": "你好，zsa！我对很多话题感兴趣，比如探索新知识和帮助解决问题。你最喜欢游泳还是篮球呢？"},
        {
            "type": "text",
            "text": "你好，zsa！我喜欢探讨各种话题和帮助解答问题。你对游泳和篮球的兴趣很广泛，有没有特别喜欢的运动方式或运动员呢？"
        },
    ]),
    HumanMessage(content="如果我想学习关于天体物理方面的知识，你能给我一些建议么？"),
    AIMessage(
        content="当然可以！你可以从基础的天文学和物理学入手，然后逐步深入到更具体的天体物理领域。阅读相关的书籍，如《宇宙的结构》或《引力的秘密》，也可以关注一些优秀的天体物理学讲座和课程。你对哪个方面最感兴趣？"
    ),
]

# 创建LLM
llm = ChatOpenAI(model="gpt-4o-mini", openai_api_base=os.getenv("OPENAI_API_BASE"))

# 使用LangGraph内置的trim_messages函数,对消息进行过滤
filtered_messages = trim_messages(
    messages=messages,  # 原始消息列表
    max_tokens=80,  # 设置最大token长度
    token_counter=llm,  # 设置LLM,用于计算token数量
    strategy="first",  # 过滤策略,可以指定从头开始或者从末尾开始
    allow_partial=True,  # 允许对消息进行分割
    text_splitter=RecursiveCharacterTextSplitter(),  # 指定文本分割器
)

print(filtered_messages)
